package main

import (
	"context"
	"flag"
	"fmt"
	"log"
	"os"
	"os/exec"
	"path/filepath"
	"strings"

	"github.com/tmc/langchaingo/llms"
	"github.com/tmc/langchaingo/llms/openai"
	"github.com/tmc/langchaingo/prompts"
	"gopkg.in/yaml.v3"
)

type CodeReviewer struct {
	llm      llms.Model
	template prompts.PromptTemplate
	allows []string
}

type Config struct {
	Token string `yaml:token`
	Model string `yaml:model`
	URL   string `yaml:url`
	Allows []string `yaml:allows`
}

func loadConfig(configPath string) (*Config, error) {
	data, err := os.ReadFile(configPath)
	if err != nil {
		return nil, err
	}
	var config Config
	err = yaml.Unmarshal(data, &config)
	if err != nil {
		return nil, err
	}
	return &config, nil
}
func NewCodeReviewer(configFile string) (*CodeReviewer, error) {
	config, configErr := loadConfig(configFile)
	if configErr != nil {
		return nil, configErr
	}

	llm, err := openai.New(
		openai.WithToken(config.Token),
		openai.WithModel(config.Model), // Specify model
		openai.WithBaseURL(config.URL), // Custom endpoint
	)
	if err != nil {
		return nil, err
	}

	template := prompts.NewPromptTemplate(`
你是专业的代码评审专家。请对这段代码进行以下方面的分析：

    漏洞与逻辑问题 ：潜在的运行时错误、空指针解引用、竞态条件
    性能 ：低效的算法、不必要的内存分配、字符串拼接问题
    风格 ：习惯用法、命名规范、错误处理模式
    安全性 ：输入验证、敏感数据处理

待评审代码：
'''
{{.code}}
'''
文件名：{{.filename}}
请提供具体、可操作的反馈。针对每个问题：

    解释为什么会成为问题
    展示如何通过代码示例进行修复
    对严重程度进行评级：严重、警告、建议

请先重点关注最重要的问题。`,
		[]string{"code", "filename"})

	return &CodeReviewer{
		llm:      llm,
		template: template,
		allows: config.Allows,
	}, nil
}

func isSupported(filename string,allows []string) bool{
	ext:=filepath.Ext(filename)
	for _,e := range allows {
		if ext == "."+e{
			return true
		}
	}
	return false
}
func (cr *CodeReviewer) ReviewFile(filename string) error {
	content, err := os.ReadFile(filename)
	if err != nil {
		return fmt.Errorf("reading file: %w", err)
	}

	if !isSupported(filename,cr.allows){
		fmt.Printf("not support this file: %s",filename)
		return nil
	}
	prompt, err := cr.template.Format(map[string]any{
		"code":     string(content),
		"filename": filename,
	})
	if err != nil {
		return fmt.Errorf("formatting prompt: %w", err)
	}

	ctx := context.Background()
	options:=[]llms.CallOption{
		llms.WithMaxTokens(40960),
		llms.WithTemperature(0.7),
	}
	response, err := cr.llm.GenerateContent(ctx, []llms.MessageContent{
		llms.TextParts(llms.ChatMessageTypeHuman, prompt),
	},options...)
	if err != nil {
		return fmt.Errorf("generating review: %w", err)
	}

	fmt.Printf("\n=== Review for %s ===\n", filename)
	reportPath := filepath.Join("report", filename)
	wrtieToFile(fmt.Sprintf("%s.md", reportPath), response.Choices[0].Content)
	fmt.Println(strings.Repeat("=", 80))
	fmt.Println(fmt.Printf("%s.md review完成", filename))
	fmt.Println(strings.Repeat("=", 80))

	return nil
}
func wrtieToFile(filename string, content string) {
	dir := filepath.Dir(filename)
	if err := os.MkdirAll(dir, 0755); err != nil {
		fmt.Println("创建目录失败: %w", err)
		return
	}
	file, err := os.Create(filename)
	if err != nil {
		fmt.Println("创建文件失败:", err)
		return
	}
	defer file.Close() // 确保文件在函数结束时关闭

	// 将内容写入文件
	_, err = file.WriteString(content)
	if err != nil {
		fmt.Println("写入文件失败:", err)
		return
	}

	fmt.Println("内容已成功写入文件:", filename)
}
func main() {
	var (
		file   = flag.String("file", "", "file to review")
		dir    = flag.String("dir", "", "Directory to review")
		config = flag.String("config", "", "config file path")
		git    = flag.Bool("git", false, "Review files changed in git working directory")
	)
	flag.Parse()

	reviewer, err := NewCodeReviewer(*config)
	if err != nil {
		log.Fatal(err)
	}

	switch {
	case *file != "":
		if err := reviewer.ReviewFile(*file); err != nil {
			log.Fatal(err)
		}
	case *dir != "":
		if err := reviewDirectory(reviewer, *dir); err != nil {
			log.Fatal(err)
		}
	case *git:
		if err := reviewGitChanges(reviewer); err != nil {
			log.Fatal(err)
		}
	default:
		fmt.Println("Usage:")
		fmt.Println("  code-reviewer -file=main.go")
		fmt.Println("  code-reviewer -dir=./pkg")
		fmt.Println("  code-reviewer -git")
		os.Exit(1)
	}
	os.Exit(1)
}

func reviewDirectory(reviewer *CodeReviewer, dir string) error {
	entris, err := os.ReadDir(dir)
	if err != nil {
		return fmt.Errorf("读取目录失败: %w", err)
	}
	for _, entry := range entris {
		path := filepath.Join(dir, entry.Name())

		if entry.IsDir() {
			err := reviewDirectory(reviewer, path)
			if err != nil {
				return err
			}
			continue
		}

		err := reviewer.ReviewFile(path)
		if err != nil {
			return err
		}
	}
	return nil
}

func reviewGitChanges(reviewer *CodeReviewer) error {
	// This is a simplified version - you'd want to use a proper git library
	cmd := exec.Command("git", "diff", "--name-only", "HEAD")
	output, err := cmd.Output()
	if err != nil {
		return fmt.Errorf("getting git changes: %w", err)
	}

	files := strings.Split(strings.TrimSpace(string(output)), "\n")
	for _, file := range files {
		if strings.HasSuffix(file, ".go") && file != "" {
			if err := reviewer.ReviewFile(file); err != nil {
				log.Printf("Error reviewing %s: %v", file, err)
			}
		}
	}
	return nil
}
